from PyQt5.Qt import *

class Edge(QGraphicsLineItem):
    def __init__(self, start_node, dest_node):
        super().__init__()
        self.start_node = start_node
        self.dest_node = dest_node
        pen = QPen(Qt.black, 2)
        pen.setCapStyle(Qt.RoundCap)
        pen.setJoinStyle(Qt.RoundJoin)
        self.setPen(pen)
        self.setFlag(QGraphicsItem.ItemIsSelectable)
        self.update_position()

    def update_position(self):
        line = self.line()
        p1x = self.start_node.scenePos().x()
        p1y = self.start_node.scenePos().y()
        width = self.start_node.rect().width()
        height = self.start_node.rect().height()
        p1 = QPointF(p1x + width/2, p1y + height)
        p2x = self.dest_node.scenePos().x()
        p2y = self.dest_node.scenePos().y()
        width_des = self.dest_node.rect().width()
        height_des = self.dest_node.rect().height()
        p2 = QPointF(p2x + width_des/2, p2y)
        line.setP1(p1)
        line.setP2(p2)
        self.setLine(line)

    def paint(self, painter, option, widget):
        # 绘制线条
        painter.setPen(self.pen())
        painter.drawLine(self.line())

        # 如果项目被选中，绘制自定义的选中标记
        if option.state & QStyle.State_Selected:
            # 设置选中标记的画笔
            painter.setPen(QPen(Qt.blue, 2, Qt.DashLine))
            # 获取线条的起点和终点
            line = self.line()
            # 绘制起点和终点的选中标记
            # p1 = line.p1()
            # p2 = line.p2()
            # painter.drawEllipse(p1, 5, 5)
            # painter.drawEllipse(p2, 5, 5)
            # 绘制线条的选中框
            painter.drawLine(line)

    def track_node(self):
        self.start_node.add_edge(self)
        self.dest_node.add_edge(self)
        self.start_node.child_node.add(self.dest_node)
        self.dest_node.parent_node.add(self.start_node)

    def remove_nodes(self):
        self.start_node.remove_edge(self)
        self.dest_node.remove_edge(self)